import os

from starlette.responses import JSONResponse

os.environ['CUDA_VISIBLE_DEVICES'] = '1'
from typing import Union, BinaryIO
import uvicorn
import numpy as np
from fastapi import FastAPI, UploadFile,status,WebSocket
from faster_whisper import WhisperModel
from time import time as ttime



from tool.logger import logger

audio_rec_save_path="./audio/received"

app = FastAPI()
model_size = "small"
# model_size = "tiny"
model = WhisperModel(model_size, device="cuda", compute_type="int8",download_root="./models",local_files_only=True)

async def exe(audio: Union[str, BinaryIO, np.ndarray],beam_size=15,language=None):
    segments, info = model.transcribe(audio, beam_size=beam_size,language=language)
    logger.info("Detected language '%s' with probability %f" % (info.language, info.language_probability))
    texts = []
    for segment in segments:
        logger.info("[%.2fs -> %.2fs] %s" % (segment.start, segment.end, segment.text))
        texts.append(segment.text)

    return " ".join(texts)



@app.post("/recognition", status_code=status.HTTP_200_OK)
async def recognition(file: UploadFile):
    start_time = ttime()
    audio=file.file
    try:
        text= await exe(audio)
        rec_time=ttime()
        # logger.info(f"保存耗时 {round(rec_time - save_time, 3)}'s")
        logger.info(f"总耗时 {round(rec_time - start_time, 3)}'s")
        return JSONResponse(
                content={'result':text}, status_code=status.HTTP_200_OK
            )
    except Exception as e:
        logger.error(e)
        return JSONResponse(
                content={'result':''}, status_code=status.HTTP_500_INTERNAL_SERVER_ERROR
            )

# @app.websocket("/recognition_stream")
# async def recognition_stream(websocket: WebSocket):
#     await websocket.accept()
#     audio_stream = bytearray()
#
#     try:
#         while True:
#             # 从WebSocket接收音频流数据
#             data = await websocket.receive_bytes()
#             # 假设客户端发送特定字符串来表示传输结束
#             if data == b"END":
#                 # 调用音频转录函数
#                 text =await exe(audio_stream)
#                 # 发送转录文本回客户端
#                 await websocket.send_text(text)
#                 # 清空音频流，为下一个传输做准备
#                 audio_stream.clear()
#             else:
#                 # 添加音频数据到流中
#                 audio_stream.extend(data)
#     except Exception as e:
#         # 在这里处理异常，如连接异常等
#         print(f"Error: {e}")
#
#     finally:
#         # 关闭WebSocket连接
#         await websocket.close()

# if __name__ == '__main__':
#     uvicorn.run(app=app, host='0.0.0.0', port=8000)